{
    "Array": 2,
    "Backend": 8,
    "Buffer": 5,
    "ClosedJaxpr": 3,
    "ConstVar": 2,
    "ConvDimensionNumbers": 3,
    "FreeVar": 1,
    "GatherDimensionNumbers": 6,
    "Jaxpr": 1,
    "JaxprEqn": 1,
    "LambdaBinding": 3,
    "Literal": 1,
    "PRNGKey": 747,
    "Partial": 1,
    "PartialVal": 6,
    "Primitive": 10,
    "RegularGridInterpolator": 4,
    "Replicated": 3,
    "ScatterDimensionNumbers": 8,
    "ShapeDtypeStruct": 15,
    "ShapedArray": 33,
    "ShardedAxis": 2,
    "ShardingSpec": 7,
    "StepTraceAnnotation": 5,
    "UndefinedPrimal": 1,
    "Unstacked": 3,
    "Var": 15,
    "Wrapped": 1,
    "__getattr__": 1,
    "_initial_style_jaxpr": 4,
    "abs": 299,
    "abstract_eval_fun": 1,
    "abstractify": 1,
    "add": 2212,
    "all": 2197,
    "all_gather": 19,
    "all_leaves": 1,
    "alltrue": 7,
    "amax": 32,
    "amin": 8,
    "any": 241,
    "apply_over_axes": 8,
    "arange": 525,
    "argmax": 109,
    "argmin": 22,
    "argwhere": 3,
    "array": 2437,
    "array_equal": 35,
    "array_split": 5,
    "as_hashable_function": 1,
    "asarray": 291,
    "atan2": 1,
    "average": 4,
    "axis_frame": 1,
    "axis_index": 4,
    "backward_pass": 10,
    "batch": 58,
    "bernoulli": 38,
    "beta": 2,
    "betainc": 1,
    "bincount": 3,
    "bitcast_convert_type": 3,
    "bitwise_and": 1,
    "bitwise_or": 1,
    "block_until_ready": 27,
    "bool_": 2,
    "broadcast": 14,
    "broadcast_arrays": 6,
    "broadcast_in_dim": 4,
    "broadcast_shapes": 33,
    "broadcast_to": 176,
    "cache": 13,
    "call_transpose": 1,
    "callback": 14,
    "canonicalize_dtype": 6,
    "categorical": 21,
    "cbrt": 1,
    "cdf": 2,
    "ceil": 74,
    "cg": 1,
    "check_type": 7,
    "checkpoint": 2,
    "cho_factor": 1,
    "cho_solve": 1,
    "choice": 81,
    "cholesky": 13,
    "choose": 1,
    "complex": 1,
    "compress": 28,
    "concatenate": 334,
    "cond": 64,
    "conj": 13,
    "constant": 239,
    "conv": 73,
    "conv_dimension_numbers": 2,
    "conv_general_dilated": 22,
    "conv_general_shape_tuple": 1,
    "conv_transpose": 7,
    "conv_transpose_shape_tuple": 1,
    "convert_element_type": 8,
    "convolve": 3,
    "convolve2d": 7,
    "copy": 435,
    "correlate": 1,
    "cos": 119,
    "cosh": 1,
    "create_token": 5,
    "cumprod": 5,
    "cumsum": 37,
    "cur_sublevel": 2,
    "custom_vjp": 4,
    "default_backend": 2,
    "defjvp": 2,
    "deflinear": 3,
    "delete": 999,
    "det": 3,
    "device_count": 83,
    "device_get": 62,
    "device_put": 27,
    "device_put_sharded": 9,
    "devices": 28,
    "diag": 51,
    "diag_indices": 2,
    "diagflat": 1,
    "digamma": 3,
    "dirichlet": 1,
    "disable_jit": 4,
    "div": 184,
    "dot": 185,
    "dot_general": 32,
    "dstack": 1,
    "dtype": 142,
    "dynamic_slice": 20,
    "dynamic_slice_in_dim": 5,
    "dynamic_update_index_in_dim": 1,
    "dynamic_update_slice": 28,
    "eig": 3,
    "eigh": 12,
    "einsum": 174,
    "einsum_path": 4,
    "empty": 11,
    "empty_like": 2,
    "entr": 3,
    "eq": 9,
    "erf": 1,
    "erfc": 1,
    "erfinv": 1,
    "eval_context": 3,
    "eval_jaxpr": 3,
    "eval_shape": 16,
    "exp": 296,
    "expand_dims": 188,
    "expit": 16,
    "expm1": 8,
    "exponential": 1,
    "extract": 33,
    "eye": 182,
    "f": 147,
    "fft": 20,
    "fftshift": 1,
    "finfo": 31,
    "flatnonzero": 1,
    "flip": 17,
    "fliplr": 2,
    "flipud": 2,
    "float32": 27,
    "floor": 49,
    "fold_in": 27,
    "fori_loop": 27,
    "frombuffer": 3,
    "fromstring": 94,
    "full": 112,
    "full_like": 15,
    "full_lower": 1,
    "gamma": 4,
    "gammainc": 1,
    "gammaincc": 2,
    "gammaln": 23,
    "gather": 11,
    "gaussian_kde": 1,
    "ge": 1,
    "gelu": 10,
    "generic": 5,
    "gensym": 1,
    "geomspace": 2,
    "get_aval": 9,
    "get_primitive_transpose": 1,
    "glorot_uniform": 4,
    "gmres": 3,
    "grad": 138,
    "gumbel": 3,
    "he_normal": 3,
    "hessian": 7,
    "hfft": 1,
    "histogram": 13,
    "host_count": 48,
    "host_id": 140,
    "host_ids": 2,
    "hsplit": 1,
    "hstack": 15,
    "i0": 1,
    "i0e": 1,
    "i1": 1,
    "i1e": 1,
    "identity": 32,
    "ifft": 7,
    "ifftshift": 1,
    "iinfo": 5,
    "imag": 4,
    "index_take": 1,
    "indices": 2,
    "infeed": 5,
    "initialize": 15,
    "insert": 144,
    "int16": 7,
    "int32": 31,
    "int64": 8,
    "integer_pow": 4,
    "inv": 19,
    "iota": 1,
    "irfft": 1,
    "is_finite": 1,
    "is_undefined_primal": 1,
    "isin": 4,
    "isneginf": 2,
    "isposinf": 2,
    "isscalar": 8,
    "issubdtype": 34,
    "ix_": 1,
    "jacfwd": 14,
    "jacobian": 15,
    "jacrev": 16,
    "jit": 209,
    "jvp": 7,
    "kaiming_normal": 3,
    "laplace": 1,
    "lecun_normal": 12,
    "lecun_uniform": 1,
    "lgamma": 1,
    "linspace": 177,
    "load": 277,
    "local_device_count": 112,
    "local_devices": 47,
    "log": 620,
    "log1p": 22,
    "logistic": 4,
    "logpdf": 3,
    "logspace": 4,
    "logsumexp": 38,
    "lower_fun": 2,
    "lpmn": 1,
    "lstsq": 1,
    "lt": 2,
    "lu": 1,
    "make_jaxpr": 13,
    "map": 416,
    "mapped_aval": 2,
    "max": 578,
    "mean": 529,
    "meshgrid": 24,
    "min": 222,
    "moveaxis": 35,
    "mul": 4,
    "multivariate_normal": 6,
    "named_call": 10,
    "namedtuple": 85,
    "ndarray": 1,
    "ndim": 49,
    "ndtr": 1,
    "ndtri": 1,
    "new_eqn_recipe": 3,
    "new_main": 3,
    "new_sublevel": 1,
    "nextafter": 1,
    "nonzero": 3,
    "normal": 435,
    "normalize": 56,
    "one_hot": 71,
    "ones": 919,
    "ones_like": 145,
    "orthogonal": 8,
    "pad": 75,
    "partial": 871,
    "pdf": 6,
    "permutation": 33,
    "pmap": 150,
    "pmax": 1,
    "pmean": 88,
    "poisson": 2,
    "polygamma": 1,
    "pow": 11,
    "power": 34,
    "ppermute": 1,
    "print": 1062,
    "process_count": 15,
    "process_index": 61,
    "prod": 115,
    "product": 35,
    "promote_types": 6,
    "psum": 57,
    "put": 72,
    "qr": 7,
    "raise_to_shaped": 5,
    "randint": 166,
    "real": 15,
    "recipe_to_eqn": 1,
    "reciprocal": 11,
    "reduce": 41,
    "reduce_window": 5,
    "ref": 1,
    "register_pytree_node": 14,
    "register_pytree_node_class": 6,
    "remat": 7,
    "repeat": 74,
    "replicate": 67,
    "reshape": 940,
    "resize": 22,
    "result_type": 34,
    "rev": 1,
    "rfft": 1,
    "rng_uniform": 3,
    "roll": 32,
    "round": 154,
    "rsqrt": 20,
    "safe_map": 24,
    "safe_zip": 6,
    "save": 3813,
    "savez": 3,
    "scan": 72,
    "scatter": 7,
    "scatter_add": 3,
    "segment_max": 5,
    "segment_sum": 22,
    "select": 172,
    "set_printoptions": 5,
    "setdiff1d": 2,
    "shape": 167,
    "shape_tensor": 1,
    "shuffle": 59,
    "shutdown": 6,
    "sign": 27,
    "sin": 111,
    "sinh": 6,
    "size": 38,
    "slice": 46,
    "slice_in_dim": 4,
    "softmax": 82,
    "solve": 20,
    "solve_triangular": 8,
    "sort": 464,
    "sort_key_val": 3,
    "sph_harm": 2,
    "split": 2605,
    "split_list": 6,
    "sqrt": 393,
    "sqrtm": 8,
    "square": 135,
    "squeeze": 140,
    "stack": 174,
    "start_trace": 1,
    "std": 63,
    "stop_gradient": 65,
    "stop_trace": 1,
    "sub": 1081,
    "subvals": 1,
    "sum": 1110,
    "svd": 6,
    "switch": 17,
    "t": 3,
    "take": 84,
    "tan": 1,
    "tanh": 53,
    "tensordot": 12,
    "tie_in": 24,
    "tile": 95,
    "top_k": 23,
    "toposort": 2,
    "trace": 45,
    "trace_to_jaxpr": 4,
    "trace_to_jaxpr_dynamic": 1,
    "transpose": 220,
    "tree_all": 1,
    "tree_flatten": 102,
    "tree_leaves": 73,
    "tree_map": 491,
    "tree_reduce": 13,
    "tree_structure": 16,
    "tree_transpose": 1,
    "tree_unflatten": 47,
    "treedef_is_leaf": 1,
    "tri": 4,
    "tril": 6,
    "tril_indices": 3,
    "tril_indices_from": 1,
    "triu": 4,
    "triu_indices": 1,
    "truncated_normal": 10,
    "uint32": 3,
    "uniform": 290,
    "unique": 48,
    "unravel_index": 1,
    "unzip2": 10,
    "value_and_grad": 97,
    "var": 36,
    "variance_scaling": 13,
    "vectorize": 4,
    "vjp": 9,
    "vmap": 284,
    "vstack": 20,
    "where": 439,
    "while_loop": 44,
    "with_sharding_constraint": 6,
    "wrap_init": 12,
    "wrap_name": 2,
    "wraps": 86,
    "xavier_normal": 13,
    "xavier_uniform": 61,
    "xla_computation": 5,
    "xlog1py": 4,
    "xlogy": 4,
    "zeros": 1060,
    "zeros_like": 179,
    "zeros_like_array": 1,
    "zeta": 1,
    "zip": 573
}